import rospy
from sensor_msgs.msg import Image
from std_msgs.msg import Bool
from geometry_msgs.msg import Point
from cv_bridge import CvBridge
from ultralytics import YOLO
import cv2
import math
import json
import torch
import os


def safe_load_model(model_path):
    "安全加载YOLO模型，处理PyTorch 2.6的权重加载问题"
    try:
        # 暂时禁用权重安全检查（仅用于可信模型）
        rospy.loginfo("Loading model with weights_only=False")
        original_load = torch.load
        def patched_load(*args, **kwargs):
            kwargs['weights_only'] = False
            return original_load(*args, **kwargs)
        torch.load = patched_load
        result = YOLO(model_path)
        torch.load = original_load  # 恢复原始函数
        return result
    except Exception as e:
        rospy.logerr(f"Failed to load model: {e}")
        raise e


class TrafficLightDetector:
    def __init__(self):
        # 初始化ROS节点
        rospy.init_node('traffic_light_detector', anonymous=False)
        
        # 创建CV Bridge用于图像转换
        self.bridge = CvBridge()
        
        # 加载YOLO模型
        model_path = os.path.join(os.path.dirname(__file__), "weights/v9 - 64 epochs.pt")
        rospy.loginfo(f"Loading model from: {model_path}")
        
        if os.path.exists(model_path):
            self.model = safe_load_model(model_path)
            rospy.loginfo("Successfully loaded custom traffic light model")
        else:
            rospy.logwarn(f"Model file not found at {model_path}, using default YOLOv8n")
            self.model = YOLO('yolov8n.pt')
        
        # 交通灯类别
        self.classNames = ["green", "red", "yellow"]
        
        # Debug变量 - True为显示图像，False为不显示
        self.debug = True
        
        # 订阅图像话题
        self.image_sub = rospy.Subscriber('/camera/image_raw', Image, self.image_callback)
        
        # 发布检测结果 - 改为Bool类型
        self.detection_pub = rospy.Publisher('/traffic_light/go_signal', Bool, queue_size=10)
        
        rospy.loginfo(f"Traffic Light Detector Node initialized with debug={self.debug}")
    
    def image_callback(self, msg):
        try:
            # 将ROS图像消息转换为OpenCV格式
            cv_image = self.bridge.imgmsg_to_cv2(msg, "bgr8")
            
            # 进行YOLO推理
            results = self.model(cv_image, stream=True)
            
            # 存储所有检测结果
            all_detections = []
            
            # 处理检测结果
            for r in results:
                boxes = r.boxes
                for box in boxes:
                    # 边界框坐标
                    x1, y1, x2, y2 = box.xyxy[0]
                    x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)
                    
                    # 置信度
                    confidence = math.ceil((box.conf[0] * 100)) / 100
                    cls = int(box.cls[0])
                    class_name = self.classNames[cls]
                    
                    # 在图像上绘制边界框
                    cv2.rectangle(cv_image, (x1, y1), (x2, y2), (255, 0, 255), 3)
                    
                    # 添加文本标签
                    label = f"{class_name}: {confidence:.2f}"
                    org = [x1, y1 - 10]
                    font = cv2.FONT_HERSHEY_SIMPLEX
                    fontScale = 0.7
                    color = (255, 0, 0)
                    thickness = 2
                    
                    cv2.putText(cv_image, label, org, font, fontScale, color, thickness)
                    
                    # 记录检测结果
                    detection = {
                        "class": class_name,
                        "confidence": float(confidence),
                        "bbox": [x1, y1, x2, y2],
                    }
                    all_detections.append(detection)
            
            # 确定最终的交通灯状态
            go_signal = self.determine_go_signal(all_detections)
            
            # 发布go_signal结果
            signal_msg = Bool()
            signal_msg.data = go_signal
            self.detection_pub.publish(signal_msg)
            
            # 如果debug模式开启，使用OpenCV显示图像
            if self.debug:
                cv2.imshow('Traffic Light Detection', cv_image)
                cv2.waitKey(1)
                
        except Exception as e:
            rospy.logerr(f"Error processing image: {e}")
    
    def determine_go_signal(self, detections):
        """
        根据检测结果确定是否可以通行
        绿灯返回True，红灯或黄灯返回False
        如果检测到多个灯，选择置信度最高的
        """
        if not detections:
            return False  # 没有检测到任何灯，默认为不可通行
        
        # 找到置信度最高的检测结果
        best_detection = max(detections, key=lambda x: x['confidence'])
        
        class_name = best_detection['class']
        confidence = best_detection['confidence']
        
        # 记录日志
        rospy.loginfo(f"Best detection: {class_name} with confidence {confidence:.2f}")
        
        # 绿灯返回True，红灯和黄灯返回False
        return class_name == "green"


def main():
    try:
        detector = TrafficLightDetector()
        rospy.spin()
    except rospy.ROSInterruptException:
        rospy.loginfo("Traffic Light Detector Node shutting down")
    finally:
        # 清理OpenCV窗口
        cv2.destroyAllWindows()


if __name__ == '__main__':
    main()